﻿using System;
using System.Diagnostics;
using System.Net;
using System.Runtime.InteropServices;

namespace RawSocketEx
{
    public class IpHelper
    {
        private const int NO_ERROR = 0;
        private const int MIB_TCP_STATE_CLOSED = 1;
        private const int MIB_TCP_STATE_LISTEN = 2;
        private const int MIB_TCP_STATE_SYN_SENT = 3;
        private const int MIB_TCP_STATE_SYN_RCVD = 4;
        private const int MIB_TCP_STATE_ESTAB = 5;
        private const int MIB_TCP_STATE_FIN_WAIT1 = 6;
        private const int MIB_TCP_STATE_FIN_WAIT2 = 7;
        private const int MIB_TCP_STATE_CLOSE_WAIT = 8;
        private const int MIB_TCP_STATE_CLOSING = 9;
        private const int MIB_TCP_STATE_LAST_ACK = 10;
        private const int MIB_TCP_STATE_TIME_WAIT = 11;
        private const int MIB_TCP_STATE_DELETE_TCB = 12;


        /*
         * Tcp Struct
         * */

        public MIB_TCPTABLE TcpConnexion;
        public MIB_EXTCPTABLE TcpExConnexions;
        public MIB_TCPSTATS TcpStats;


        /*
         * Udp Struct
         * */
        public MIB_UDPTABLE UdpConnexion;
        public MIB_EXUDPTABLE UdpExConnexion;
        public MIB_UDPSTATS UdpStats;

        #region Tcp Function

        public void GetTcpStats()
        {
            TcpStats = new MIB_TCPSTATS();
            IPHlpAPI32Wrapper.GetTcpStatistics(ref TcpStats);
        }


        public void GetExTcpConnexions()
        {
            // the size of the MIB_EXTCPROW struct =  6*DWORD
            const int rowsize = 24;
            int BufferSize = 100000;
            // allocate a dumb memory space in order to retrieve  nb of connexion
            IntPtr lpTable = Marshal.AllocHGlobal(BufferSize);
            //getting infos
            int res = IPHlpAPI32Wrapper.AllocateAndGetTcpExTableFromStack(ref lpTable, true,
                                                                          IPHlpAPI32Wrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                Debug.WriteLine("Erreur : " + IPHlpAPI32Wrapper.GetAPIErrorMessageDescription(res) + " " + res);
                return; // Error. You should handle it
            }
            int CurrentIndex = 0;
            //get the number of entries in the table
            var NumEntries = (int)Marshal.ReadIntPtr(lpTable);
            // free allocated space in memory
            Marshal.FreeHGlobal(lpTable);


            ///////////////////
            // calculate the real buffer size nb of entrie * size of the struct for each entrie(24) + the dwNumEntries
            BufferSize = (NumEntries * rowsize) + 4;
            // make the struct to hold the resullts
            TcpExConnexions = new MIB_EXTCPTABLE();
            // Allocate memory
            lpTable = Marshal.AllocHGlobal(BufferSize);
            res = IPHlpAPI32Wrapper.AllocateAndGetTcpExTableFromStack(ref lpTable, true,
                                                                      IPHlpAPI32Wrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                Debug.WriteLine("Erreur : " + IPHlpAPI32Wrapper.GetAPIErrorMessageDescription(res) + " " + res);
                return; // Error. You should handle it
            }
            // New pointer of iterating throught the data
            IntPtr current = lpTable;

            CurrentIndex = 0;
            // get the (again) the number of entries
            NumEntries = (int)Marshal.ReadIntPtr(current);
            TcpExConnexions.dwNumEntries = NumEntries;
            // Make the array of entries
            TcpExConnexions.table = new MIB_EXTCPROW[NumEntries];
            // iterate the pointer of 4 (the size of the DWORD dwNumEntries)
            CurrentIndex += 4;
            current = (IntPtr)((int)current + CurrentIndex);
            // for each entries
            for (int i = 0; i < NumEntries; i++)
            {
                // The state of the connexion (in string)
                TcpExConnexions.table[i].StrgState = ConvertState((int)Marshal.ReadIntPtr(current));
                // The state of the connexion (in ID)
                TcpExConnexions.table[i].iState = (int)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                // get the local address of the connexion
                var localAddr = (UInt32)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                // get the local port of the connexion
                var localPort = (UInt32)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                // Store the local endpoint in the struct and convertthe port in decimal (ie convert_Port())
                TcpExConnexions.table[i].Local = new IPEndPoint(localAddr, ConvertPort(localPort));
                // get the remote address of the connexion
                var RemoteAddr = (UInt32)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                UInt32 RemotePort = 0;
                // if the remote address = 0 (0.0.0.0) the remote port is always 0
                // else get the remote port
                if (RemoteAddr != 0)
                {
                    RemotePort = (UInt32)Marshal.ReadIntPtr(current);
                    RemotePort = ConvertPort(RemotePort);
                }
                current = (IntPtr)((int)current + 4);
                // store the remote endpoint in the struct  and convertthe port in decimal (ie convert_Port())
                TcpExConnexions.table[i].Remote = new IPEndPoint(RemoteAddr, (int)RemotePort);
                // store the process ID
                TcpExConnexions.table[i].dwProcessId = (int)Marshal.ReadIntPtr(current);
                // Store and get the process name in the struct
                TcpExConnexions.table[i].ProcessName = GetProcessName(TcpExConnexions.table[i].dwProcessId);
                current = (IntPtr)((int)current + 4);
            }
            // free the buffer
            Marshal.FreeHGlobal(lpTable);
        }


        public void GetTcpConnexions()
        {
            var buffer = new byte[20000]; // Start with 20.000 bytes left for information about tcp table
            int pdwSize = 20000;
            int res = IPHlpAPI32Wrapper.GetTcpTable(buffer, out pdwSize, true);
            if (res != NO_ERROR)
            {
                buffer = new byte[pdwSize];
                res = IPHlpAPI32Wrapper.GetTcpTable(buffer, out pdwSize, true);
                if (res != 0)
                    return; // Error. You should handle it
            }

            TcpConnexion = new MIB_TCPTABLE();

            int nOffset = 0;
            // number of entry in the
            TcpConnexion.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            nOffset += 4;
            TcpConnexion.table = new MIB_TCPROW[TcpConnexion.dwNumEntries];

            for (int i = 0; i < TcpConnexion.dwNumEntries; i++)
            {
                // state
                int st = Convert.ToInt32(buffer[nOffset]);
                // state in string
                ((TcpConnexion.table[i])).StrgState = ConvertState(st);
                // state  by ID
                ((TcpConnexion.table[i])).iState = st;
                nOffset += 4;
                // local address
                string LocalAdrr = buffer[nOffset] + "." + buffer[nOffset + 1] + "." + buffer[nOffset + 2] + "." +
                                   buffer[nOffset + 3];
                nOffset += 4;
                //local port in decimal
                int LocalPort = ((buffer[nOffset]) << 8) + ((buffer[nOffset + 1])) +
                                ((buffer[nOffset + 2]) << 24) + ((buffer[nOffset + 3]) << 16);

                nOffset += 4;
                // store the remote endpoint
                ((TcpConnexion.table[i])).Local = new IPEndPoint(IPAddress.Parse(LocalAdrr), LocalPort);

                // remote address
                string RemoteAdrr = buffer[nOffset] + "." + buffer[nOffset + 1] + "." + buffer[nOffset + 2] + "." +
                                    buffer[nOffset + 3];
                nOffset += 4;
                // if the remote address = 0 (0.0.0.0) the remote port is always 0
                // else get the remote port in decimal
                int RemotePort;
                //
                if (RemoteAdrr == "0.0.0.0")
                {
                    RemotePort = 0;
                }
                else
                {
                    RemotePort = ((buffer[nOffset]) << 8) + ((buffer[nOffset + 1])) +
                                 ((buffer[nOffset + 2]) << 24) + ((buffer[nOffset + 3]) << 16);
                }
                nOffset += 4;
                ((TcpConnexion.table[i])).Remote = new IPEndPoint(IPAddress.Parse(RemoteAdrr), RemotePort);
            }
        }

        #endregion

        #region Udp Functions

        public void GetUdpStats()
        {
            UdpStats = new MIB_UDPSTATS();
            IPHlpAPI32Wrapper.GetUdpStatistics(ref UdpStats);
        }

        public void GetUdpConnexions()
        {
            byte[] buffer = null;
            int pdwSize;
            var res = IPHlpAPI32Wrapper.GetUdpTable(null, out pdwSize, true);

            if (res != NO_ERROR)
            {
                buffer = new byte[pdwSize];
                res = IPHlpAPI32Wrapper.GetUdpTable(buffer, out pdwSize, true);
                if (res != 0)
                    return;
            }

            if (buffer == null) return;

            UdpConnexion = new MIB_UDPTABLE();

            var nOffset = 0;
            UdpConnexion.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            UdpConnexion.table = new MIB_UDPROW[UdpConnexion.dwNumEntries];
            nOffset += 4;
            for (int i = 0; i < UdpConnexion.dwNumEntries; i++)
            {
                string LocalAdrr = string.Format("{0}.{1}.{2}.{3}", buffer[nOffset], buffer[nOffset + 1], buffer[nOffset + 2], buffer[nOffset + 3]);
                nOffset += 4;

                int LocalPort = ((buffer[nOffset]) << 8) + ((buffer[nOffset + 1])) +
                                ((buffer[nOffset + 2]) << 24) + ((buffer[nOffset + 3]) << 16);
                nOffset += 4;

                ((UdpConnexion.table[i])).Local = new IPEndPoint(IPAddress.Parse(LocalAdrr), LocalPort);
            }
        }

        public void GetExUdpConnexions()
        {
            // the size of the MIB_EXTCPROW struct =  4*DWORD
            int rowsize = 12;
            int BufferSize = 100000;
            // allocate a dumb memory space in order to retrieve  nb of connexion
            IntPtr lpTable = Marshal.AllocHGlobal(BufferSize);
            //getting infos
            int res = IPHlpAPI32Wrapper.AllocateAndGetUdpExTableFromStack(ref lpTable, true,
                                                                          IPHlpAPI32Wrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                Debug.WriteLine("Erreur : " + IPHlpAPI32Wrapper.GetAPIErrorMessageDescription(res) + " " + res);
                return; // Error. You should handle it
            }
            int CurrentIndex = 0;
            //get the number of entries in the table
            var NumEntries = (int)Marshal.ReadIntPtr(lpTable);
            lpTable = IntPtr.Zero;
            // free allocated space in memory
            Marshal.FreeHGlobal(lpTable);

            ///////////////////
            // calculate the real buffer size nb of entrie * size of the struct for each entrie(24) + the dwNumEntries
            BufferSize = (NumEntries * rowsize) + 4;
            // make the struct to hold the resullts
            UdpExConnexion = new MIB_EXUDPTABLE();
            // Allocate memory
            lpTable = Marshal.AllocHGlobal(BufferSize);
            res = IPHlpAPI32Wrapper.AllocateAndGetUdpExTableFromStack(ref lpTable, true,
                                                                      IPHlpAPI32Wrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                Debug.WriteLine("Erreur : " + IPHlpAPI32Wrapper.GetAPIErrorMessageDescription(res) + " " + res);
                return; // Error. You should handle it
            }
            // New pointer of iterating throught the data
            IntPtr current = lpTable;
            CurrentIndex = 0;
            // get the (again) the number of entries
            NumEntries = (int)Marshal.ReadIntPtr(current);
            UdpExConnexion.dwNumEntries = NumEntries;
            // Make the array of entries
            UdpExConnexion.table = new MIB_EXUDPROW[NumEntries];
            // iterate the pointer of 4 (the size of the DWORD dwNumEntries)
            CurrentIndex += 4;
            current = (IntPtr)((int)current + CurrentIndex);
            // for each entries
            for (int i = 0; i < NumEntries; i++)
            {
                // get the local address of the connexion
                var localAddr = (UInt32)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                // get the local port of the connexion
                var localPort = (UInt32)Marshal.ReadIntPtr(current);
                // iterate the pointer of 4
                current = (IntPtr)((int)current + 4);
                // Store the local endpoint in the struct and convertthe port in decimal (ie convert_Port())
                UdpExConnexion.table[i].Local = new IPEndPoint(localAddr, ConvertPort(localPort));
                // store the process ID
                UdpExConnexion.table[i].dwProcessId = (int)Marshal.ReadIntPtr(current);
                // Store and get the process name in the struct
                UdpExConnexion.table[i].ProcessName = GetProcessName(UdpExConnexion.table[i].dwProcessId);
                current = (IntPtr)((int)current + 4);
            }
            // free the buffer
            Marshal.FreeHGlobal(lpTable);
            // re init the pointer
            current = IntPtr.Zero;
        }

        #endregion

        #region helper fct

        private static UInt16 ConvertPort(UInt32 dwPort)
        {
            var b = new Byte[2];
            // high weight byte
            b[0] = byte.Parse((dwPort >> 8).ToString());
            // low weight byte
            b[1] = byte.Parse((dwPort & 0xFF).ToString());
            return BitConverter.ToUInt16(b, 0);
        }

        private static string ConvertState(int state)
        {
            string strgState = "";
            switch (state)
            {
                case MIB_TCP_STATE_CLOSED:
                    strgState = "CLOSED";
                    break;
                case MIB_TCP_STATE_LISTEN:
                    strgState = "LISTEN";
                    break;
                case MIB_TCP_STATE_SYN_SENT:
                    strgState = "SYN_SENT";
                    break;
                case MIB_TCP_STATE_SYN_RCVD:
                    strgState = "SYN_RCVD";
                    break;
                case MIB_TCP_STATE_ESTAB:
                    strgState = "ESTAB";
                    break;
                case MIB_TCP_STATE_FIN_WAIT1:
                    strgState = "FIN_WAIT1";
                    break;
                case MIB_TCP_STATE_FIN_WAIT2:
                    strgState = "FIN_WAIT2";
                    break;
                case MIB_TCP_STATE_CLOSE_WAIT:
                    strgState = "CLOSE_WAIT";
                    break;
                case MIB_TCP_STATE_CLOSING:
                    strgState = "CLOSING";
                    break;
                case MIB_TCP_STATE_LAST_ACK:
                    strgState = "LAST_ACK";
                    break;
                case MIB_TCP_STATE_TIME_WAIT:
                    strgState = "TIME_WAIT";
                    break;
                case MIB_TCP_STATE_DELETE_TCB:
                    strgState = "DELETE_TCB";
                    break;
            }
            return strgState;
        }


        private static string GetProcessName(int processId)
        {
            //could be an error here if the process die before we can get his name
            try
            {
                Process p = Process.GetProcessById(processId);
                return p.ProcessName;
            }
            catch (Exception ex)
            {
                Debug.WriteLine(ex.Message);
                return "unknown";
            }
        }

        #endregion

        public MIB_TCPROW_OWNER_PID[] GetAllTcpConnections()
        {
            MIB_TCPROW_OWNER_PID[] tTable;
            const int AF_INET = 2;
            int buffSize = 0;

            uint ret = IPHlpAPI32Wrapper.GetExtendedTcpTable(IntPtr.Zero, ref buffSize, true, AF_INET,
                                                             TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL, 0);
            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);

            try
            {
                ret = IPHlpAPI32Wrapper.GetExtendedTcpTable(buffTable, ref buffSize, true, AF_INET,
                                                            TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL, 0);
                if (ret != 0)
                {
                    return null;
                }

                var tab = (MIB_TCPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_TCPTABLE_OWNER_PID));
                var rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(tab.dwNumEntries));
                tTable = new MIB_TCPROW_OWNER_PID[tab.dwNumEntries];

                for (int i = 0; i < tab.dwNumEntries; i++)
                {
                    var tcpRow = (MIB_TCPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_TCPROW_OWNER_PID));
                    tTable[i] = tcpRow;
                    rowPtr = (IntPtr)((long)rowPtr + Marshal.SizeOf(tcpRow));
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffTable);
            }


            return tTable;
        }

        public MIB_UDPROW_OWNER_PID[] GetAllUdpConnections()
        {
            MIB_UDPROW_OWNER_PID[] tTable;
            const int afInet = 2;
            int buffSize = 0;

            uint ret = IPHlpAPI32Wrapper.GetExtendedUdpTable(IntPtr.Zero, ref buffSize, true, afInet,
                                                             UDP_TABLE_CLASS.UDP_TABLE_OWNER_PID, 0);

            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);

            try
            {
                ret = IPHlpAPI32Wrapper.GetExtendedUdpTable(buffTable, ref buffSize, true, afInet,
                                                            UDP_TABLE_CLASS.UDP_TABLE_OWNER_PID, 0);
                if (ret != 0)
                {
                    return null;
                }

                var tab = (MIB_UDPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_UDPTABLE_OWNER_PID));
                var rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(tab.dwNumEntries));
                tTable = new MIB_UDPROW_OWNER_PID[tab.dwNumEntries];

                for (int i = 0; i < tab.dwNumEntries; i++)
                {
                    var udpRow = (MIB_UDPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_UDPROW_OWNER_PID));
                    tTable[i] = udpRow;
                    rowPtr = (IntPtr)((long)rowPtr + Marshal.SizeOf(udpRow)); // next entry
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffTable);
            }
            return tTable;
        }
    }
}